import numpy as np
import scipy.spatial.distance
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from mykmeans import MyKMeans
from utils import JS_D
import time
import pickle

training_round = 2
sync_round = 2



candidate1_round = list(set([1,2,3]) - set([sync_round]))[0]
candidate2_round = list(set([1,2,3]) - set([sync_round]))[1]


tar_subtitle = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round%d/round%d_withPreLabel_scpByTypicalScene.srt'%(training_round, sync_round)
scp_path = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round%d/scp_byTypicalScene.npy'%training_round


src_subtitle = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round%d/round%d.srt'%(sync_round, sync_round)
sync_ds_path       = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round%d/round%d_%dds.npy'%(training_round, training_round, sync_round)
candidate1_ds_path = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round%d/round%d_%dds.npy'%(training_round, training_round, candidate1_round)
candidate2_ds_path = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round%d/round%d_%dds.npy'%(training_round, training_round, candidate2_round)

index_map_file = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/IndexMap.pickle'

f = open(src_subtitle,'r')
src_lines = f.readlines()
f.close()

scp = np.load(scp_path)

sync_ds = np.load(sync_ds_path)
candidate1_ds = np.load(candidate1_ds_path)
candidate2_ds = np.load(candidate2_ds_path)


frame_count = np.shape(sync_ds)[0]
sync_pre_labels = np.zeros(frame_count, dtype=int)
candidate1_pre_labels = np.zeros(frame_count, dtype=int)
candidate2_pre_labels = np.zeros(frame_count, dtype=int)

index_map_file = open(index_map_file,'rb')
index_map = pickle.load(index_map_file)
index_map_file.close()
mapping = index_map[sync_round - 1]

print('Processing...')

for i in range(frame_count):
    i1 = mapping[i][candidate1_round-1]
    i2 = mapping[i][candidate2_round-1]

    jsd_min = np.inf
    for j in range(3):
        jsd = JS_D(sync_ds[i], scp[j])
        if jsd < jsd_min:
            jsd_min = jsd
            sync_pre_labels[i] = j
            
    jsd_min = np.inf
    for j in range(3):
        jsd = JS_D(candidate1_ds[i1], scp[j])
        if jsd < jsd_min:
            jsd_min = jsd
            candidate1_pre_labels[i] = j

    jsd_min = np.inf
    for j in range(3):
        jsd = JS_D(candidate2_ds[i2], scp[j])
        if jsd < jsd_min:
            jsd_min = jsd
            candidate2_pre_labels[i] = j


label_text = ['Straight road','Trun road', 'Alerting traffic']

f = open(tar_subtitle,'w')
fno = 0
for line in src_lines:
    if line[0] == 'f':
        f.write(line)
        f.write('round %d:%sround %d:%sround %d:%s\n'%(sync_round, label_text[sync_pre_labels[fno]].center(20), candidate1_round, label_text[candidate1_pre_labels[fno]].center(20), candidate2_round, label_text[candidate2_pre_labels[fno]].rjust(20)))
        fno += 1
    else:
        f.write(line)
f.close()

print('Program exit normally.')